import torch
from torchvision import datasets
from torchvision.transforms import transforms
import torch.utils.data


class DataHandler:
    def __init__(self, data_dir, batch_size):
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.transform = transforms.Compose([
            transforms.Resize((300, 300)),
            transforms.RandomHorizontalFlip(),  # 这是一个随机水平翻转操作，有一定概率（通常是50%）将图像水平翻转
            transforms.RandomVerticalFlip(),  # 随机垂直翻转操作
            transforms.RandomAutocontrast(),  # 随机自动对比度调整操作
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])

    def preprocess_data(self):
        train_dataset = datasets.ImageFolder(self.data_dir, transform=self.transform)
        train_size = int(0.7 * len(train_dataset))
        val_size = len(train_dataset) - train_size
        train_dataset, val_dataset = torch.utils.data.random_split(train_dataset, [train_size, val_size])
        return train_dataset, val_dataset

    def preprocess_data_loader(self):
        train_dataset = datasets.ImageFolder(self.data_dir, transform=self.transform)
        train_size = int(0.7 * len(train_dataset))
        val_size = len(train_dataset) - train_size
        train_dataset, val_dataset = torch.utils.data.random_split(train_dataset, [train_size, val_size])

        train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=self.batch_size, shuffle=True)
        val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=self.batch_size)

        return train_loader, val_loader


class Trainer:
    def __init__(self, model, device, loss_function, optimizer):
        self.model = model
        self.device = device
        self.loss_function = loss_function
        self.optimizer = optimizer

    def train(self, dataloader):
        self.model.train()  # 设置模型为训练模式
        size = len(dataloader.dataset)

        for batch, (images, labels) in enumerate(dataloader):
            images, labels = images.to(self.device), labels.to(self.device)

            # 前向传播
            pred = self.model(images)
            # 计算预测和损失
            loss = self.loss_function(pred, labels)

            # 反向传播和优化
            loss.backward()
            self.optimizer.step()
            self.optimizer.zero_grad()
            print('>', end='')
            # 每批中第10（9+1）个打印信息
            if batch % 10 == 9:
                loss, current = loss.item(), (batch + 1) * len(images)
                print(f"Training loss: {loss:>7f}  [{batch + 1} * len(X) ] [{current}/{size}]")

    def test(self, dataloader):
        self.model.eval()  # 设置模型为评估模式
        size = len(dataloader.dataset)
        num_batches = len(dataloader)
        test_loss, correct = 0, 0

        with torch.no_grad():
            for X, y in dataloader:
                X, y = X.to(self.device), y.to(self.device)
                pred = self.model(X)
                test_loss += self.loss_function(pred, y).item()
                correct += (pred.argmax(1) == y).type(torch.float).sum().item()

        test_loss /= num_batches
        correct /= size

        print(f"Test Error: Accuracy: {(100 * correct):.1f}%, Avg loss: {test_loss:.8f}")
